# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import functools
import os
import ujson as json
from collections import namedtuple

from coin.base.datetime_util import to_datetime
from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.bitmex.book_builder import BookBuilder
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.bitmex.liquidation_book_builder import (LiquidationBookBuilder, new_order_entry)
from coin.exchange.base.ws.trade_aggregator import TradeAggregator

BitmexInstrument = namedtuple(
    'BitmexInstrument',
    ['timestamp', 'market_type', 'exchange', 'native_symbol', 'inst', 'topic_id'])


class BitmexHandler(WebsocketHandler):
  def __init__(self, *, event_dist=None, feed_checker_config=None, agg_view=True):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._trade_processor = TradeAggregator(handler=self, agg_view=agg_view)
    self._instruments = {}
    self._books = {}
    self._liquidation_books = {}

  def on_ws_orderbook_l2(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    updated_books = {}
    action = message['action']
    if 'table' in message and message['table'] == 'orderBookL2':
      if action == 'update':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            book.update(data['id'], data['side'], data['size'])
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'insert':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            book.insert(data)
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'delete':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._books:
            book = self._books[symbol]
            book.delete(data['id'], data['side'])
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'partial':
        assert message['keys'] == ['symbol', 'id', 'side']
        reset_symbols = set()
        for data in message['data']:
          symbol = data['symbol']
          product = BitmexFuturesProduct.FromStrNativeProductNothrow(symbol)
          if product is None:
            continue

          if not self.is_subscribed('book', symbol):
            continue

          if symbol not in self._books:
            self._books[symbol] = BookBuilder(symbol,
                                              book_builder_observer=self._book_builder_observer)
            reset_symbols.add(symbol)

          book = self._books[symbol]
          if symbol not in reset_symbols:
            book.clear_book()
            reset_symbols.add(symbol)

          book.partial(data)
          book.update_timestamp(record.timestamp)
          updated_books[symbol] = book
        self.ready = True

      else:
        raise ValueError('Unknown action: %s' % action)

    for book in updated_books.values():
      valid = False
      try:
        # Raise attribute error if bid0/ask0 is None
        valid = book.bid0().price < book.ask0().price
      except AttributeError:
        pass

      if valid:
        book.update_topic_id(record.topic_id)
        self.publish_event('book', book.symbol, book)

      if self._feed_checker:
        self._feed_checker.on_book(product=book.product, book=book, timestamp=record.timestamp)

  def on_ws_trade(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    if message['action'] == 'partial':
      return

    self._trade_processor.on_begin_trade()

    assert message['action'] == 'insert', message
    for trade_msg in message['data']:
      symbol = trade_msg['symbol']
      product = BitmexFuturesProduct.FromStrNativeProductNothrow(symbol)
      if product is None:
        continue
      if not self.is_subscribed('trade', symbol):
        continue

      price = float(trade_msg['price'])
      qty = float(trade_msg['size'])

      if trade_msg['side'] == 'Buy':
        side = TradeSide.TRADE_BUY_SIDE
      elif trade_msg['side'] == 'Sell':
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError('Unknown side: %s' % trade_msg['side'])

      trade_event = create_market_trade_event(record.timestamp,
                                              symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      self._trade_processor.on_mid_trade(trade_event)

      if symbol[0] != '.':  # Skip index product.
        if self._feed_checker:
          self._feed_checker.on_trade(product=product,
                                      trade_event=trade_event,
                                      timestamp=record.timestamp)

    self._trade_processor.on_end_trade()

  def on_ws_instrument(self, record, queue_data, topic_data):
    """Handling instruments. Care when you subscribe new symbol."""
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    assert message['table'] == 'instrument'

    action = message['action']
    if action == 'partial':
      assert message['keys'] == ['symbol']

      self._instruments = {}
      for data in message['data']:
        symbol = data['symbol']
        if self.is_subscribed('instrument', symbol):
          assert symbol not in self._instruments

          self._instruments[symbol] = dict(data)
          self.publish_event(
              'instrument',
              symbol,
              BitmexInstrument(record.timestamp,
                               'Futures',
                               'Bitmex',
                               symbol,
                               self._instruments[symbol],
                               record.topic_id))

    elif action == 'insert':
      # Replace if there is existing data else append it.
      for data in message['data']:
        symbol = data['symbol']
        if self.is_subscribed('instrument', symbol):
          self._instruments[symbol] = dict(data)
          self.publish_event(
              'instrument',
              symbol,
              BitmexInstrument(record.timestamp,
                               'Futures',
                               'Bitmex',
                               symbol,
                               self._instruments[symbol],
                               record.topic_id))

    elif action == 'delete':
      for data in message['data']:
        symbol = data['symbol']
        if self.is_subscribed('instrument', symbol):
          self._instruments.pop(symbol, None)

    elif action == 'update':
      updated_symbols = set()
      for data in message['data']:
        symbol = data['symbol']
        if (not self.is_subscribed('instrument', symbol) or symbol not in self._instruments):
          continue

        self._instruments[symbol].update(data)
        updated_symbols.add(symbol)

      # Publish merged updates.
      for symbol in updated_symbols:
        self.publish_event(
            'instrument',
            symbol,
            BitmexInstrument(record.timestamp,
                             'Futures',
                             'Bitmex',
                             symbol,
                             self._instruments[symbol],
                             record.topic_id))

    else:
      raise ValueError('Unknown action: %s' % action)

  def subscribe_instrument(self, symbols, callback):
    self.subscribe_event('instrument', symbols, callback)

  def on_ws_liquidation(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    updated_books = {}
    action = message['action']
    if 'table' in message and message['table'] == 'liquidation':
      if action == 'partial':
        self._liquidation_books.clear()
        assert message['keys'] == ['orderID']
        self._liquidation_books = {}
        for data in message['data']:
          symbol = data['symbol']
          if not self.is_subscribed('liquidation', symbol):
            continue
          if symbol not in self._liquidation_books:
            self._liquidation_books[symbol] = LiquidationBookBuilder(symbol)
          book = self._liquidation_books[symbol]
          book.insert(new_order_entry(data))
          book.update_timestamp(record.timestamp)
          updated_books[symbol] = book

      elif action == 'delete':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._liquidation_books:
            book = self._liquidation_books[symbol]
            book.delete(data['orderID'])
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'insert':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._liquidation_books:
            book = self._liquidation_books[symbol]
            book.insert(new_order_entry(data))
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book
          elif self.is_subscribed('liquidation', symbol):
            self._liquidation_books[symbol] = LiquidationBookBuilder(symbol)
            book = self._liquidation_books[symbol]
            book.insert(new_order_entry(data))
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      elif action == 'update':
        for data in message['data']:
          symbol = data['symbol']
          if symbol in self._liquidation_books:
            book = self._liquidation_books[symbol]
            book.update(data['orderID'], new_order_entry(data))
            book.update_timestamp(record.timestamp)
            updated_books[symbol] = book

      else:
        raise ValueError('Unknown action: %s' % action)

    for book in updated_books.values():
      try:
        self.publish_event('liquidation', book.symbol, book)
      except IndexError:
        continue

  def subscribe_liquidation(self, symbols, callback):
    self.subscribe_event('liquidation', symbols, callback)


if ('RUST_FEED_HANDLER' in os.environ) and os.environ['RUST_FEED_HANDLER'] == '1':
  try:
    from native_rust import feed_handler

    class BitmexRustHandler(BitmexHandler):
      def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._rust_handler = feed_handler.BitmexHandler()
        self._rust_handler.register_book_callback(functools.partial(self.publish_event, 'book'))

      def post_subscribe(self, event_type, symbol, callback):
        assert type(symbol) == str
        if event_type == 'book':
          self._rust_handler.register_book_symbol(symbol)

      def on_ws_orderbook_l2(self, record, queue_data, topic_data):
        if isinstance(record.data, (dict, list)):
          record_bytes = json.dumps(record.data).encode()
        else:
          record_bytes = record.data.tobytes()

        self._rust_handler.on_orderbook_l2(record.timestamp, record_bytes, {}, {})

    BitmexHandler = BitmexRustHandler

  except ImportError:
    import logging
    logging.warning('Unable to import rust feed_handler.')
